#include "HetmansProblem.h"

using namespace std;

HetmansProblem::HetmansProblem ( int n ) : N ( n ), MAX_WEIGHT ( combination ( N ) ) {
    state_.resize ( N, 0 );
    calculateCurrentEnergy();
}

HetmansProblem::HetmansProblem ( const std::vector< int >& start_state ) : 
    N ( start_state.size() ), 
    MAX_WEIGHT ( combination ( N ) ) {
    state_ = start_state;
    calculateCurrentEnergy();
}

HetmansProblem::~HetmansProblem() {

}


float HetmansProblem::weight ( int state ) const  {
    vector<int> stateVect = state_;
    stateVect.at ( state / N ) = state % N;
    return calculateWeight ( stateVect );
}

void HetmansProblem::setState ( int state ) {
    state_.at ( state / N ) = state % N;
    calculateCurrentEnergy();
}

void HetmansProblem::generateStates ( std::vector< int >& container ) {
    container.clear();
    for ( int col = 0; col < N; ++col ) {
        for ( int row = 0; row < N; ++row ) {
            if ( state_.at ( col ) == row ) {
                continue;
            }
            container.push_back ( N * col + row );
        }
    }
}

void HetmansProblem::calculateCurrentEnergy() {
    currentWeight_ = calculateWeight ( state_ );
}

float HetmansProblem::calculateWeight ( const std::vector< int >& state ) const {
    float energy = 0;
    vector<int> conflicts ( N, 0 );

    for ( int i = 0; i < N; ++i ) {
        ++conflicts.at ( state.at ( i ) );
    }

    for ( int i = 0; i < N; ++i ) {
        if ( conflicts.at ( i ) > 1 ) {
            energy += combination ( conflicts.at ( i ) );
        }
    }

    conflicts.clear();
    conflicts.resize ( 2 * N - 1, 0 );

    for ( int i = 0; i < N; ++i ) {
        int idx = i - state.at ( i ) + N - 1;
        ++conflicts.at ( idx );
    }

    for ( int i = 0; i <  2 * N - 1; ++i ) {
        if ( conflicts.at ( i ) > 1 ) {
            energy += combination ( conflicts.at ( i ) );
        }
    }

    conflicts.clear();
    conflicts.resize ( 2 * N - 1, 0 );

    for ( int i = 0; i < N; ++i ) {
        int idx = state.at ( i ) + i;
        ++conflicts.at ( idx );
    }


    for ( int i = 0; i < 2 * N - 1; ++i ) {
        if ( conflicts.at ( i ) > 1 ) {
            energy += combination ( conflicts.at ( i ) );
        }
    }

    return energy;
}

bool HetmansProblem::isSolved() const {
    return currentWeight_ == 0;
}

float HetmansProblem::currentEnergy() const {
    return MAX_WEIGHT - currentWeight_;
}

float HetmansProblem::energy ( int state ) const {
    return MAX_WEIGHT - weight ( state );
}
